package ace.cmp.spring.graphql.test.test.core.graphql;

import ace.cmp.spring.graphql.test.test.core.model.Author;
import ace.cmp.spring.graphql.test.test.core.model.Book;
import ace.cmp.spring.graphql.test.test.core.repository.AuthorRepository;
import ace.cmp.spring.graphql.test.test.core.repository.BookRepository;
import java.util.HashMap;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import lombok.extern.slf4j.Slf4j;
import org.dataloader.DataLoader;
import org.springframework.graphql.data.method.annotation.*;
import org.springframework.graphql.execution.BatchLoaderRegistry;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Mono;

/**
 * @author caspar
 * @date 2023/8/22 17:08
 */
@Slf4j
@RestController
public class BookController {

  public static final String QUERY_TYPE_NAME = "BookNameSpace";
  public static final String MUTATION_TYPE_NAME = "BookNameSpace";
  private final BookRepository bookRepository;
  private final AuthorRepository authorRepository;
  private final BatchLoaderRegistry batchLoaderRegistry;

  public BookController(
      BookRepository bookRepository,
      AuthorRepository authorRepository,
      BatchLoaderRegistry batchLoaderRegistry) {
    this.bookRepository = bookRepository;
    this.authorRepository = authorRepository;
    this.batchLoaderRegistry = batchLoaderRegistry;
    batchLoaderRegistry
        .forTypePair(Long.class, Book.class)
        .registerMappedBatchLoader(
            (ids, env) -> {
              log.info("getBookLoadById:" + ids.toString());
              return Mono.fromSupplier(
                  () -> {
                    List<Book> books = this.bookRepository.getByIdIn(ids.stream().toList());
                    return ids.stream()
                        .collect(
                            HashMap::new,
                            (m, id) -> {
                              Book book =
                                  books.stream()
                                      .filter(p2 -> id.equals(p2.getId()))
                                      .findFirst()
                                      .orElse(null);
                              m.put(id, book);
                            },
                            HashMap::putAll);
                  });
            });
    batchLoaderRegistry
        .forTypePair(Long.class, Author.class)
        .registerMappedBatchLoader(
            (ids, env) -> {
              log.info("getAuthorLoadById:" + ids.toString());
              return Mono.fromSupplier(
                  () -> {
                    List<Author> authors = this.authorRepository.getByIdIn(ids.stream().toList());
                    return ids.stream()
                        .collect(
                            HashMap::new,
                            (m, id) -> {
                              Author book =
                                  authors.stream()
                                      .filter(p2 -> id.equals(p2.getId()))
                                      .findFirst()
                                      .orElse(null);
                              m.put(id, book);
                            },
                            HashMap::putAll);
                  });
            });
  }

  @QueryMapping(value = "books")
  public BookController books() {
    return this;
  }

  @SchemaMapping(typeName = BookController.QUERY_TYPE_NAME, field = "getBookById")
  public CompletableFuture<Book> getBookById(@Argument Long id, DataLoader<Long, Book> loader) {

    log.info("请求：getBookById");
    return loader.load(id);
  }

  @SchemaMapping(typeName = BookController.QUERY_TYPE_NAME, field = "getAllBooks")
  public List<Book> getAllBooks() {

    return bookRepository.getAll();
  }

  @QueryMapping("getAllBooks")
  public List<Book> getAllBooksByQuery() {

    return bookRepository.getAll();
  }

  @SchemaMapping(typeName = BookController.QUERY_TYPE_NAME, field = "modifyBookTitleById")
  public Boolean modifyBookTitleById(@Argument("id") Long id, @Argument("title") String title) {
    bookRepository.getById(id).setTitle(title);
    return true;
  }

  @QueryMapping("getBookById")
  public CompletableFuture<Book> getBookById2(@Argument Long id, DataLoader<Long, Book> loader) {

    log.info("请求getBookById: " + id.toString());
    return loader.load(id);
  }

  /**
   * 关联查询对象，不能使用该形式实现，会有N+1查询问题
   */
  @SchemaMapping(typeName = "Book", field = "authorN")
  public Author getAuthorN(Book book) {
    log.info("请求getAuthorN：" + book.getId().toString());
    return this.authorRepository.getById(book.getId());
  }

  /**
   * 查询批量关联对象，并解决N+1查询问题
   *
   * @param book
   * @return
   */
  @SchemaMapping(typeName = "Book", field = "author")
  public CompletableFuture<Author> getAuthor(Book book, DataLoader<Long, Author> loader) {
    log.info("请求getAuthorByAuthorId ：" + book.getAuthorId().toString());
    return loader.load(book.getAuthorId());
  }
}
